# coding:utf-8

import os
import shutil


def batch_split_txt_into_chapters(source_dir, output_root):
    """
    批量处理文件夹中的txt文件，按分隔符拆分章节并保存

    参数:
        source_dir: 存放原始txt文件的文件夹路径
        output_root: 保存拆分后章节的根目录路径
    """
    # 确保输出根目录存在
    if not os.path.exists(output_root):
        os.makedirs(output_root)

    # 遍历源目录中的所有文件
    for filename in os.listdir(source_dir):
        # 只处理txt文件
        if filename.endswith('.txt'):
            # 获取书名（不含扩展名）
            book_name = os.path.splitext(filename)[0]
            # 创建该书的输出文件夹
            book_output_dir = os.path.join(output_root, book_name)

            # 如果文件夹已存在，先清空
            if os.path.exists(book_output_dir):
                shutil.rmtree(book_output_dir)
            os.makedirs(book_output_dir)

            # 读取文件内容
            file_path = os.path.join(source_dir, filename)
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    content = f.read()

                # 按分隔符拆分章节（过滤空内容）
                chapters = [chap.strip() for chap in content.split('###>>>') if chap.strip()]

                # 保存每个章节
                for i, chapter in enumerate(chapters, 1):
                    chapter_filename = f"{book_name}_{i}.txt"
                    chapter_path = os.path.join(book_output_dir, chapter_filename)

                    with open(chapter_path, 'w', encoding='utf-8') as f:
                        f.write(chapter)

                print(f"处理完成：《{book_name}》，共拆分出 {len(chapters)} 个章节")

            except Exception as e:
                print(f"处理《{book_name}》时出错：{str(e)}")


if __name__ == "__main__":
    # 配置路径（可根据实际情况修改）
    source_directory = "./books"  # 存放原始txt文件的文件夹
    output_directory = "./split_books"  # 拆分后章节的保存目录

    # 执行拆分
    batch_split_txt_into_chapters(source_directory, output_directory)
    print("所有文件处理完毕！")
